package midware

import (
	"github.com/cihub/seelog"
	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
	"shop/config/constant"
	"shop/utils/result"
	"shop/utils/sessionutil"
	"strconv"
)

// limitType
const accessLimitIp = 1
const accessLimitApi = 2
const accessLimitUserId = 3

func IpLimiter(putRate, bucketSize int) func(*gin.Context) {
	return Limiter(accessLimitIp, putRate, bucketSize)
}

func ApiLimiter(putRate, bucketSize int) func(*gin.Context) {
	return Limiter(accessLimitApi, putRate, bucketSize)
}

func UserIdLimiter(putRate, bucketSize int) func(*gin.Context) {
	return Limiter(accessLimitUserId, putRate, bucketSize)
}

var limiterCache = map[string]*rate.Limiter{} // <key,*rate.Limiter>

// Limiter 限流，令牌桶算法， putRate:每秒放置令牌的数量，bucketSize:令牌桶的大小
func Limiter(limitType, putRate, bucketSize int) func(*gin.Context) {
	return func(ctx *gin.Context) {
		// seelog.Debugf("accessLimiter, limitType=%d, putRate=%d, bucketSize=%d", limitType, putRate, bucketSize)
		key := ""
		if limitType == accessLimitIp {
			key = constant.CacheKeyPrefixAccessLimiter + ctx.ClientIP()
		} else if limitType == accessLimitApi {
			key = constant.CacheKeyPrefixAccessLimiter + ctx.FullPath()
		} else if limitType == accessLimitUserId {
			key = constant.CacheKeyPrefixAccessLimiter + strconv.FormatInt(sessionutil.GetUserIdFromCtx(ctx), 10)
		} else {
			seelog.Warnf("invalid limitType:%d", limitType)
			ctx.Next()
			return
		}

		if limiter, ok := limiterCache[key]; ok {
			if limiter.Allow() {
				ctx.Next()
			} else {
				seelog.Debugf("访问频繁，limitType=%d, path=%s", limitType, ctx.FullPath())
				result.RenderResult(ctx, result.Fail("访问频繁，请稍后重试"))
				ctx.Abort()
			}
		} else {
			limiterCache[key] = rate.NewLimiter(rate.Limit(putRate), bucketSize) // warn:非原子操作
			ctx.Next()
		}
	}
}
